from renderer import TextureSpaceRender

import torch
import os
import matplotlib.pyplot as plt
import numpy as np
import mitsuba as mi
import config as C
import pytorch_config as PC
from util import to_db_number_type, debug, info

mi.set_variant("scalar_rgb")


def load_diffuse_normal_image(index: int) -> torch.DoubleTensor:
    p_image = plt.imread(C.resource_base_path + "diffNormals/" + to_db_number_type(index) + ".png")
    p_image = p_image * 2 - 1
    p_image /= np.expand_dims(np.linalg.norm(p_image, 2, axis=2), axis=-1)
    tensor = torch.DoubleTensor(p_image)
    return tensor


def load_diffuse_albedo_image(index: int) -> torch.DoubleTensor:
    p_image = plt.imread(C.resource_base_path + "diffAlbedo/" + to_db_number_type(index) + ".png")
    tensor = torch.DoubleTensor(p_image)
    return tensor


def load_specular_normal_image(index: int) -> torch.DoubleTensor:
    p_image = plt.imread(C.resource_base_path + "specNormals/" + to_db_number_type(index) + ".png")
    p_image = p_image * 2 - 1
    p_image /= np.expand_dims(np.linalg.norm(p_image, 2, axis=2), axis=-1)
    tensor = torch.DoubleTensor(p_image)
    return tensor


def load_specular_albedo_image(index: int) -> torch.DoubleTensor:
    p_image = plt.imread(C.resource_base_path + "specAlbedo/" + to_db_number_type(index) + ".png")
    tensor = torch.DoubleTensor(p_image)
    tensor = tensor.unsqueeze(dim=-1).expand(-1, -1, 3)
    return tensor


def load_image_in_batch(batch_size, bias=0) -> (torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor):
    if batch_size <= 0:
        ...
    diff_normal_list = []
    diff_albedo_list = []
    spec_normal_list = []
    spec_albedo_list = []
    for i in range(batch_size):
        diff_normal_list.append(load_diffuse_normal_image(index=i + bias))
        diff_albedo_list.append(load_diffuse_albedo_image(index=i + bias))
        spec_normal_list.append(load_specular_normal_image(index=i + bias))
        spec_albedo_list.append(load_specular_albedo_image(index=i + bias))
    batch_diff_normal = torch.stack(diff_normal_list, dim=0)
    batch_diff_albedo = torch.stack(diff_albedo_list, dim=0)
    batch_spec_normal = torch.stack(spec_normal_list, dim=0)
    batch_spec_albedo = torch.stack(spec_albedo_list, dim=0)
    return batch_diff_albedo, batch_diff_normal, batch_spec_albedo, batch_spec_normal


def batch_output(result, output_path: str, bias=0):
    batch_size = result.shape[0]
    for i in range(batch_size):
        # using mitsuba's util to write image async.
        mi.util.write_bitmap(output_path + f"{bias + i}.png", np.array(result[i]))


def main():
    """
        procedure entrance.
    """
    '''
        Initialize setting.
    '''
    total_size = 1500
    batch_size = 50
    texture_renderer = TextureSpaceRender(
        light_direction=torch.DoubleTensor(PC.light_direction),
        light_radiance=torch.DoubleTensor(PC.light_radiance) if PC.using_hdr else torch.DoubleTensor([1.0, 1.0, 1.0]),
        view_direction=torch.DoubleTensor(PC.view_direction),
        using_sss=PC.using_sss,
        using_gamma_correction=PC.using_gamma_correction,
        using_specular_normal_when_diffuse=PC.using_specular_normal_when_diffuse,
        using_hdr=PC.using_hdr
    )
    # Ensure filepath exists.
    if not os.path.exists(PC.output_path):
        os.makedirs(PC.output_path, exist_ok=True)
    info(f"output file is : {PC.output_path}")

    '''
        Separate total batch into several turns.
    '''
    render_turns = total_size // batch_size
    bias = 0
    for _ in range(render_turns):
        info(f"Start render images {bias} -> {bias + batch_size - 1}")
        '''
                Load image and concatenate them into batch.
        '''
        info("Start loading images.")
        diff_albedo, diff_normal, spec_albedo, spec_normal = load_image_in_batch(batch_size, bias=bias)
        if C.DEBUG:
            debug(
                f"DA :{diff_albedo.shape}\tDN :{diff_normal.shape}\tSA :{spec_albedo.shape}\tSN :{spec_normal.shape}")
        info("Images loading finished.")

        '''
            Render
        '''
        info("Start rendering.")
        result = texture_renderer(
            diff_albedo,
            diff_normal,
            spec_albedo,
            spec_normal
        )
        if C.DEBUG:
            debug(f"result shape\t:{result.shape}")
        info("Rendering finished.")

        '''
            Output batch result.
        '''
        batch_output(result, PC.output_path, bias=bias)
        bias += batch_size


if __name__ == "__main__":
    main()
